Add ToolsRetriever class and Retriever.convert_to_tool() method #332

oskarhane · 2025-05-08T10:47:06Z

Description

In addition to the major parts of this PR, I've refactored and moved the tool.py file into its own directory to group relevant tool files together. This move touches a lot of files because of updated imports.

`ToolsRetriever`

Now that we have tool support in the LLM interface, let's create a retriever that can use one or more tools to retrieve data.

This retriever follows the Retriever interface so it can be used within the GraphRAG class to get the full e2e experience (see example file).

The ToolsRetriever uses an LLM to decide on what tools to use to find the relevant data.

calendar_tool = CalendarTool()
weather_tool = WeatherTool()

llm = OpenAILLM(...)
driver = GraphDatabase.driver(...)

tools_retriever = ToolsRetriever(
    driver=driver,
    llm=llm,
    tools=[calendar_tool, weather_tool],
)

graphrag = GraphRAG(
    llm=llm,
    retriever=tools_retriever,
)

result = graphrag.search(query_text="Tell me about tomorrow", return_context=False)

`convert_to_tool ()`

This method is a way to convert a Retriever to a Tool so it can be used within the ToolsRetriever. This is useful when you might want to have both a VectorRetriever and a Text2CypherRetreiver as a fallback.
See new example files for usage.

Type of Change

How Has This Been Tested?

Unit tests
E2E tests
Manual tests

Checklist

The following requirements should have been met (depending on the changes in the branch):

Documentation has been updated
Unit tests have been updated
E2E tests have been updated
Examples have been updated
New files have copyright header
CLA (https://neo4j.com/developer/cla/) has been signed
CHANGELOG.md updated if appropriate

examples/retrieve/tools/retriever_to_tool_example.py

stellasia · 2025-05-13T14:12:53Z

examples/retrieve/tools/retriever_to_tool_example.py

+    )
+
+    # Convert the retriever to a tool with specific parameters
+    static_tool = convert_retriever_to_tool(


I'd say this should be a method from the Retriever class, and parameters should be encapsulated in this class as well, these parameters are bound to the search method and won't change from one instance to another.

So something like:

class Retriever: def get_parameters(self) -> ObjectParameter: raise NotImplementedError() # need to be implemented in subclasses def convert_to_tool(self, name: str, description: Optional[str] = None) -> Tool: # rest of the function goes here

Note: as a future improvement, I think we could infer the parameters from the search method signature without having to redeclare it.

Yeah, makes sense. I wanted to be as little intrusive as possible, bot I've updated to a method like this now.

stellasia · 2025-05-13T14:16:46Z

src/neo4j_graphrag/retrievers/tools_retriever.py

+
+    def __init__(
+        self,
+        driver: neo4j.Driver,


Should we consider updating the Retriever interface if this one does not use the driver?

I'm not sure tbh. I think the tools retriever is an exception so I think we can keep it for now.

src/neo4j_graphrag/retrievers/tools_retriever.py

stellasia · 2025-05-13T14:21:50Z

src/neo4j_graphrag/retrievers/tools_retriever.py

+                    # Execute the tool with the provided arguments
+                    tool_result = selected_tool.execute(**tool_args)
+                    # If the tool result is a RawSearchResult, extract its records
+                    if hasattr(tool_result, "records"):


We can maybe enforce the fact that Tool.execute must return a list of something (or list of strings?), or do we really want to create fake records?

Yeah, I'll think about this one a bit.

src/neo4j_graphrag/tools/tool.py

src/neo4j_graphrag/tools/utils.py

stellasia · 2025-05-13T14:26:42Z

src/neo4j_graphrag/tools/utils.py

+        # arguments like query_text, top_k, etc., passed as keyword arguments.
+        # The Tool's 'parameters' definition (e.g., ObjectParameter) ensures
+        # that these arguments are provided in kwargs when Tool.execute is called.
+        return retriever.get_search_results(**kwargs)


Why not using the search method?

I think dealing with the raw output is better in these cases, and let the ToolRetriever handle the formatting.

Each retriever is likely to return results in different format, I mean it will "often" be a neo4j Record, but with different keys, especially if using the "*CypherRetriever", so the only formatting we can do here is stringify these records, right? But we can let the formatting question for a later discussion, it's probably part of a wider decision.

- Add abstract get_parameters() method to Retriever base class - Add convert_to_tool() instance method to Retriever class - Implement get_parameters() for all concrete retriever classes - Remove automatic query_text injection in ToolsRetriever - Update example to use new convert_to_tool() method - Remove unnecessary description from ObjectParameter in example

stellasia · 2025-06-19T12:10:07Z

src/neo4j_graphrag/retrievers/base.py

+        """
+        return self._infer_parameters_from_signature(parameter_descriptions or {})
+
+    def _infer_parameters_from_signature(


src/neo4j_graphrag/retrievers/base.py

stellasia · 2025-06-19T12:21:12Z

src/neo4j_graphrag/tools/tool.py

            self._parameters = parameters
-
-        self._execute_func = execute_func
+        else:


Would it be worth checking parameters is really None here? The user could provide a different type (a list of something for instance) and silently not using it would lead to unexpected behavior

stellasia · 2025-06-19T12:23:20Z

src/neo4j_graphrag/retrievers/tools_retriever.py

+        self,
+        driver: neo4j.Driver,
+        llm: LLMInterface,
+        tools: Sequence[Tool],


Are we checking somewhere that tool's name are unique? Or this is not a requirement for the LLM?

stellasia · 2025-06-19T12:27:14Z

src/neo4j_graphrag/tools/utils.py

+from neo4j_graphrag.tools.tool import Tool, ObjectParameter
+
+
+def convert_retriever_to_tool(


Isn't this function replaced by Retriever.convert_to_tool?

oskarhane requested a review from a team as a code owner May 8, 2025 10:47

stellasia reviewed May 13, 2025

View reviewed changes

Add ToolsRetriever class and convert_retriever_to_tool() fn

d7a0104

oskarhane force-pushed the tools-retriever-3 branch from 63f3348 to 3eb99e5 Compare June 18, 2025 15:45

oskarhane changed the title ~~Add ToolsRetriever class and convert_retriever_to_tool() function~~ Add ToolsRetriever class and Retriever.convert_to_tool() method Jun 18, 2025

oskarhane force-pushed the tools-retriever-3 branch from 3eb99e5 to b4e532e Compare June 18, 2025 20:52

oskarhane requested a review from stellasia June 18, 2025 21:01

stellasia reviewed Jun 19, 2025

View reviewed changes

src/neo4j_graphrag/retrievers/base.py Show resolved Hide resolved

stellasia reviewed Jun 19, 2025

View reviewed changes

		from neo4j_graphrag.tools.tool import Tool, ObjectParameter


		def convert_retriever_to_tool(

Add ToolsRetriever class and Retriever.convert_to_tool() method #332

Are you sure you want to change the base?

Add ToolsRetriever class and Retriever.convert_to_tool() method #332

Uh oh!

Conversation

oskarhane commented May 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

ToolsRetriever

convert_to_tool ()

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

Uh oh!

stellasia May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stellasia May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

stellasia May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stellasia Jun 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

oskarhane commented May 8, 2025 •

edited

Loading

`ToolsRetriever`

`convert_to_tool ()`

stellasia May 13, 2025 •

edited

Loading

stellasia May 13, 2025 •

edited

Loading

stellasia May 13, 2025 •

edited

Loading

stellasia Jun 19, 2025 •

edited

Loading